package SecureUtil;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.KeyFactory;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.KeyPair;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.spec.EncodedKeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.Cipher;

public class RSAHelper {

    public PrivateKey privKey;
    public PublicKey pubKey;

    /**
     * 
     * Used to generate keys and save them to the public parameters
     * 
     * @throws IOException
     * @throws NoSuchProviderException
     */
    public void generateKeys() throws IOException, NoSuchProviderException {
        String algorithm = "RSA";
        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(algorithm);
            SecureRandom random = SecureRandom.getInstance("SHA1PRNG", "SUN");
            keyGen.initialize(1024, random);
            KeyPair keyPair = keyGen.generateKeyPair();

            privKey = keyPair.getPrivate();
            pubKey = keyPair.getPublic();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("usage: java AsymmetricKeyMaker <RSA | DSA>");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * IMPORTANT: if you want to use other's public key to encrypt the data,
     * 			  you need to assign the publicKey FIRST.
     * 
     * Use public Key to encrypt data,
     * 
     * @param in
     * @return
     * @throws Exception
     */
    public byte[] RSAEncrypt(byte[] in) throws Exception {
        Cipher pubCipher = Cipher.getInstance("RSA");
        pubCipher.init(Cipher.ENCRYPT_MODE, pubKey);
        /*
        byte[] block = new byte[117];
        int j = 0;
        for(int i=0; i<in.length; i++){
            block[j] = in[i];
            j++;
            if(j==117){
                if(i==in.length-1) return pubCipher.doFinal(block);
                else pubCipher.update(block);
                j=0;
            }
        }
        
        if(j!=0){
            byte[] lastblock = null;
            lastblock = new byte[j];
            for(int i=0; i< j; i++)
                lastblock[i]=block[j];
            return pubCipher.doFinal(lastblock);
        }*/
        return pubCipher.doFinal(in);
        //throw new Exception("exception occurs in breaking to block for encryption");
    }

    
    /**
     * If you want to decrypt the data, make sure you use the correct Key
     * 
     * @param in
     * @return
     * @throws Exception
     */
    public byte[] RSADecrypt(byte[] in) throws Exception {
        Cipher prvCipher = Cipher.getInstance("RSA");
        prvCipher.init(Cipher.DECRYPT_MODE, privKey);
        /*
        byte[] block = new byte[117];
        int j = 0;
        for(int i=0; i<in.length; i++){
            block[j] = in[i];
            j++;
            if(j==117){
                if(i==in.length-1) return prvCipher.doFinal(block);
                else prvCipher.update(block);
                j=0;
            }
        }
        
        if(j!=0){
            byte[] lastblock = null;
            lastblock = new byte[j];
            for(int i=0; i< j; i++)
                lastblock[i]=block[j];
                //return prvCipher.doFinal(lastblock);
        }
         */
        return prvCipher.doFinal(in);
        //throw new Exception("exception occurs in breaking to block for decryption");
    }

    /**
     * Signature, make sure you use the correct private key to sign
     * 
     * @param in
     * @return
     * @throws Exception
     */
    public byte[] RSASignature(byte[] in) throws Exception {
        Signature sig = Signature.getInstance("SHA512withRSA");
        sig.initSign(privKey);
        sig.update(in);
        return sig.sign();
    }

    /**
     * verify
     * 
     * 1st parameter is original data, 2nd parameter is the signature that you 
     * receive from others
     * 
     * @param tobeverified
     * @param sig
     * @return
     * @throws Exception
     */
    public boolean RSASigVerify(byte[] tobeverified, byte[] sig) throws Exception {
        Signature siger = Signature.getInstance("SHA512withRSA");
        siger.initVerify(pubKey);
        siger.update(tobeverified);
        return siger.verify(sig);
    }

    
    /**
     * First call generateKeys(), then use this method to save keys to file
     */
    public void storeKeys() {
        try {
            byte[] pubKeyBytes = pubKey.getEncoded();
            byte[] privKeyBytes = privKey.getEncoded();

            FileOutputStream privkeyfos = new FileOutputStream("private_key");
            privkeyfos.write(privKeyBytes);
            privkeyfos.close();

            FileOutputStream pubkeyfos = new FileOutputStream("public_key");
            pubkeyfos.write(pubKeyBytes);
            pubkeyfos.close();
        } catch (FileNotFoundException e) {
            System.err.println("File operation error");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * This method read keys from files
     * 
     * @throws Exception
     */
    public void restoreKeys() throws Exception {
        privKey = getPrivateKeyFromFile("private_key");
        pubKey = getPublicKeyFromFile("public_key");
    }

    /**
     * Read private key from a certain file name
     * 
     * @param keyfile
     * @return
     * @throws Exception
     */
    public PrivateKey getPrivateKeyFromFile(String keyfile) throws Exception {
        FileInputStream in = new FileInputStream(keyfile);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        byte[] buf = new byte[1024 * 4];
        in.read(buf);
        EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(buf);
        PrivateKey privateKey = keyFactory.generatePrivate(privateKeySpec);
        in.close();
        return privateKey;
    }

    
    /**
     * Read public key from a certain file name
     * 
     * @param keyfile
     * @return
     * @throws Exception
     */
    public PublicKey getPublicKeyFromFile(String keyfile) throws Exception {
        FileInputStream in = new FileInputStream(keyfile);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        byte[] buf = new byte[1024 * 4];
        in.read(buf);
        EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(buf);
        PublicKey publicKey = keyFactory.generatePublic(publicKeySpec);
        in.close();
        return publicKey;
    }
}
